#include <stdint.h>
#include <signal.h>
#include "NameService.h"
#include "SocketUtil.h"
#ifndef WIN32
#include <sys/time.h>
#endif
//uint16_t NameService::nameTransId_ = 0x3131;
uint32_t NameService::myTTL_ = 0;

NameService::NameService()
{
    char hostname[256] = {'\0'};
    //name_t localhostname;

    NbtUtil::init();

    int res = gethostname(hostname, sizeof(hostname));
    if (res != 0)
        Debug(2, "get the host name failed\n");
	else 
		Debug(1, "HOST NAME IS:%s\n", hostname);
    
    myname_ = string(hostname);

    /* 
    localhostname.name = myname_; 
    localhostname.ttl = 0;
    localhostname.nb_flags = 0x80; //why??
    localhostname.valid = true;
    names_.push_back(localhostname);
    */
    struct hostent *hp;
    struct sockaddr_in sock;
    hp = gethostbyname(myname_.c_str());
    if (hp == NULL)
    {
        Debug(0, "gethostbyname failed\n");
        return ;
    }

    memset(&sock, 0, sizeof(sock));
    memcpy(&sock.sin_addr, hp->h_addr, hp->h_length);
    sock.sin_port = htons(myNetBiosPort);
    sock.sin_family = hp->h_addrtype;
    sock.sin_addr.s_addr = INADDR_ANY;

    this->sockfd_ = socket(hp->h_addrtype, SOCK_DGRAM, 0);
    if (this->sockfd_ == INVALID_SOCKET)
    {
        perror("create a listen socket failed\n");
        return;
    }
    //allow broadcast
    int one=1;
    setsockopt(sockfd_,SOL_SOCKET,SO_BROADCAST,&one,sizeof(one));

    //setsockopt();
    /* now we've got a socket - we need to bind it */
    if (bind(sockfd_, (struct sockaddr*)&sock,sizeof(sock)) < 0) 
    {
        Debug(2, "bind the socket failed\n");
        closesocket(sockfd_);
        exit(0) ;
    }

#ifndef WIN32
    /* We will abort gracefully when the client or remote system 
       goes away */
    signal(SIGPIPE, SIG_IGN);
#endif

    struct sockaddr_in localip, brcastip;

    NbtUtil::get_ipaddr(&localip, &brcastip);
    myip_ = localip.sin_addr;
    bcastIp_ = brcastip.sin_addr;
    nameTransId_ = 0x3121;
    Debug(2,"Local Ip:%s, broadcast Ip:%s\n", inet_ntoa(myip_), inet_ntoa(bcastIp_));
    /*
       bool ret = registerName(localhostname.name, &myip_);
       if (!ret)
       names_[0].valid = false;
       */

    /* the code below can't work to get the local ip */
    /*
       struct hostent *h_ptr;
       struct sockaddr_in myip_addr;
       h_ptr = gethostbyname2(hostname, AF_INET);
       if (h_ptr == NULL)
       Debug(2, "get the hostent by hostname %s failed\n", hostname);
       else 
       {
       int k = 0;
       while (h_ptr->h_addr_list[k] != NULL)
       {

       }
       }  */
}

bool NameService::registerName(string name, const struct in_addr *ip)
{
  unsigned char nb_flags = 0;
  struct in_addr lastip;
  int count;
  char *p;

  size_t k;
  for (k = 0; k < registerednames_.size(); ++k)
  {
	  if (registerednames_[k] == name)
		  break;
  }
  if (k >= registerednames_.size())
	  registerednames_.push_back(name);

  Debug(1,"Registering name %s (%s) nb_flags=0x%x\n",
	name.c_str(), inet_ntoa(*ip) ,nb_flags);

  SSVAL(outBuffer_,0,++nameTransId_); 
  //CVAL(outBuffer_,2) = (0x5<<3) | 0x1;
  SCVAL(outBuffer_, 2, ((0x5 << 3) | 0x1));
  SCVAL(outBuffer_,3, ((1<<4) | 0x0));
  SSVAL(outBuffer_,4,1);
  SSVAL(outBuffer_,6,0);
  SSVAL(outBuffer_,8,0);
  SSVAL(outBuffer_,10,1);  
  p = outBuffer_+12;
  NbtUtil::name_mangle(name.c_str(),p);
  p += NbtUtil::name_len(p);
  SSVAL(p,0,0x20);
  SSVAL(p,2,0x01);
  p += 4;
  SCVAL(p,0,0xC0);
  SCVAL(p,1,12);
  p += 2;
  SSVAL(p,0,0x20);
  SSVAL(p,2,0x1);
  SIVAL(p,4,myTTL_);
  SSVAL(p,8,6);
  SCVAL(p,10,nb_flags);
  SCVAL(p,11,0);
  p += 12;
  memcpy(p,ip,4);
  p += 4;

  count = 3;
  while (count--)
  {
      Debug(2,"Sending register request for %s at (%s)\n",
              name.c_str(), inet_ntoa(*ip));

      NbtUtil::show_nmb(outBuffer_);
      if (!NbtUtil::send_packet(outBuffer_,
           NbtUtil::nmb_len(outBuffer_),&bcastIp_,myNetBiosPort,SOCK_DGRAM))
	        return false;
receiveStage:
      if (NbtUtil::receive_nmb(sockfd_, inBuffer_,
          sizeof(inBuffer_), &lastip, 3))
	  {
          if (lastip.s_addr == myip_.s_addr)
          {
              goto receiveStage;
          }
          int rec_name_trn_id = SVAL(inBuffer_,0);
          int opcode = CVAL(inBuffer_,2) >> 3;
          int nm_flags = ((CVAL(inBuffer_,2) & 0x7) << 4) + (CVAL(inBuffer_,3)>>4);
          int rcode = CVAL(inBuffer_,3) & 0xF;

          /* is it a positive response to our request? */
          if ((rec_name_trn_id = nameTransId_) && 
                  (opcode == 5 && nm_flags == 0x58 && rcode != 0 && rcode != 7))
         {
              char qname[100];
              NbtUtil::name_extract(inBuffer_,12,qname);
              if (NbtUtil::name_equal(qname,name.c_str()))
             {
              Debug(0,"Someone (%s) gave us a negative name regregistration response!\n",
                inet_ntoa(lastip));
              return false;
             }
	      }	  
	  
          /* it's something else - process it anyway, unless we are running
           as a daemon. This is necessary as we may have been started by a 
           name query of our name arriving on port myNetBiosPort (often happens) */
          //if (!daemon)
          //{
              NbtUtil::show_nmb(inBuffer_);
              reply(inBuffer_,outBuffer_ + NbtUtil::nmb_len(outBuffer_), lastip);
          //}
     }
  }
  Debug(1, "no negative responce\n");
  /* no negative replies, send a demand */
  p = outBuffer_;
  SSVAL(outBuffer_,0,++nameTransId_);
  SCVAL(outBuffer_,2,(0x5<<3));
  SCVAL(outBuffer_,3,(1<<4) | 0x0);
  SSVAL(outBuffer_,4,1);
  SSVAL(outBuffer_,6,0);
  SSVAL(outBuffer_,8,0);
  SSVAL(outBuffer_,10,1);  
  p = outBuffer_+12;
  NbtUtil::name_mangle(name.c_str(),p);
  p += NbtUtil::name_len(p);
  SSVAL(p,0,0x20);
  SSVAL(p,2,0x1);
  p += 4;
  SCVAL(p,0,0xC0);
  SCVAL(p,1,12);
  p += 2;
  SSVAL(p,0,0x20);
  SSVAL(p,2,0x1);
  SIVAL(p,4,myTTL_);
  SSVAL(p,8,6);
  SCVAL(p,10,nb_flags);
  SCVAL(p,11,0);
  p += 12;
  memcpy(p,ip,4);
  p += 4;
  
  Debug(2,"Sending reg demand for %s at (%s)\n",
          name.c_str(),inet_ntoa(*ip));

  NbtUtil::show_nmb(outBuffer_);
  return(NbtUtil::send_packet(outBuffer_,
              NbtUtil::nmb_len(outBuffer_),
              &bcastIp_,myNetBiosPort,SOCK_DGRAM));
}

bool NameService::queryName(string name, struct in_addr* ip)
{
  unsigned char nb_flags = 0;
  struct in_addr lastip;
  int count;
  int wait_time = 5;
  char *p;

  Debug(1,"Querying name %s nb_flags=0x%x\n",
	name.c_str(), nb_flags);

  SSVAL(outBuffer_,0,++nameTransId_);
  SCVAL(outBuffer_,2,0x01);
  SCVAL(outBuffer_,3,0x10);
  SSVAL(outBuffer_,4,1);
  SSVAL(outBuffer_,6,0);
  SSVAL(outBuffer_,8,0);
  SSVAL(outBuffer_,10,0);  
  p = outBuffer_+12;
  NbtUtil::name_mangle(name.c_str(),p);
  p += NbtUtil::name_len(p);
  SSVAL(p,0,0x20);
  SSVAL(p,2,0x1);
  p += 4;

  count = 3;
  while (count--)
  {
      Debug(2,"Sending query request for %s at (%s)\n",
              name.c_str(),inet_ntoa(myip_));


      NbtUtil::show_nmb(outBuffer_);
      if (!NbtUtil::send_packet(outBuffer_,
           NbtUtil::nmb_len(outBuffer_),&bcastIp_,
           myNetBiosPort,SOCK_DGRAM))
	        return false;
      if (count == 1)
          wait_time = 20;
receiveStage:
      if (NbtUtil::receive_nmb(sockfd_, inBuffer_, 
            sizeof(inBuffer_), &lastip, wait_time))
	  {
          if (lastip.s_addr == myip_.s_addr)
          {
              goto receiveStage;
          }
          int rec_name_trn_id = SVAL(inBuffer_,0);
          int opcode = CVAL(inBuffer_,2) >> 3;
          int nm_flags = ((CVAL(inBuffer_,2) & 0x7) << 4) + (CVAL(inBuffer_,3)>>4);
          int rcode = CVAL(inBuffer_,3) & 0xF;

          /* is it a negative response to our request? */
          if ((rec_name_trn_id = nameTransId_) && 
                  (opcode == 0 && (nm_flags&~0x28) == 0x50 && rcode == 0))
          {
              char qname[100];
              NbtUtil::name_extract(inBuffer_,12,qname);
              if (NbtUtil::name_equal(qname,name.c_str()))
              {
                  Debug(0,"Someone (%s) gave us a positive name query response!\n",
                    inet_ntoa(lastip));
                  char *ptr = inBuffer_ + 12;
                  int  namelen = NbtUtil::name_len(ptr);
                  ptr += (namelen + 12);
                  memcpy(ip, ptr, sizeof(*ip));
                  return true;
              }
          }	  
          else if ((rec_name_trn_id = nameTransId_) && 
                   (opcode == 0 && (nm_flags&~0x08) == 0x50 && rcode != 0))
          {
            return false;
          }
          
          /* it's something else - process it anyway, unless we are running
           as a daemon. This is necessary as we may have been started by a 
           name query of our name arriving on port myNetBiosPort (often happens) */
          //if (!daemon)
           // {
          NbtUtil::show_nmb(inBuffer_);
          reply(inBuffer_,outBuffer_ + NbtUtil::nmb_len(outBuffer_), lastip);
            //}
	  }
    }
    return false;

}

//
void NameService::replyRegDemand(char *inbuf, char *outbuf, const struct in_addr lastip)
{
	int name_trn_id = SVAL(inbuf,0);
	char qname[100]="";
	int ttl;
	char *p = inbuf;
	struct in_addr ip;
	int n=0;
	unsigned char nb_flags;

	NbtUtil::name_extract(inbuf,12,qname);
	p += 12;
	p += NbtUtil::name_len(p);
	p += 4;
	p += NbtUtil::name_len(p);
	p += 4;
	ttl = IVAL(p,0);
	nb_flags = CVAL(p,6);
	p += 8;
	memcpy(&ip,p,4);

	Debug(2,"Name registration overwrite and demand for %s (%s) nb_flags=0x%x\n",
		qname,inet_ntoa(ip),nb_flags);
	/* if the name doesn't exist yet then don't respond */
	if ((n = findName(qname)) < 0) 
	{
		addName(qname, ip);
		return;
	}
}

//
void NameService::replyRegRequest(char *inbuf, char *outbuf, const struct in_addr lastip)
{
  int name_trn_id = SVAL(inbuf,0);
  char qname[100]="";
  int ttl;
  char *p = inbuf;
  struct in_addr ip;
  int n=0;
  unsigned char nb_flags;
  
  NbtUtil::name_extract(inbuf,12,qname);

  p += 12;
  p += NbtUtil::name_len(p);
  p += 4;
  p += NbtUtil::name_len(p);
  p += 4;
  ttl = IVAL(p,0);
  nb_flags = CVAL(p,6);
  p += 8;
  memcpy(&ip,p,4);

  Debug(2,"Name registration request for %s (%s) nb_flags=0x%x\n",
	qname,inet_ntoa(ip),nb_flags);

//  /* if it's a group name and being registered as a group then it's OK */
//  if ((names_[n].nb_flags & 0x80) && (nb_flags & 0x80))
//    return;

  /* if it's not my name then don't worry about it, because 
   * there must be another one defend against it */
  if (!NbtUtil::name_equal(myname_.c_str(),qname))
    return;

  Debug(0,"Someones using my name (%s), sending negative reply\n",qname);

  /* Send a NEGATIVE REGISTRATION RESPONSE to protect our name */
  SSVAL(outbuf,0,name_trn_id);
  SCVAL(outbuf,2,(1<<7) | (0x5<<3) | 0x5);
  SCVAL(outbuf,3,(1<<7) | 0x6);
  SSVAL(outbuf,4,0);
  SSVAL(outbuf,6,1);
  SSVAL(outbuf,8,0);
  SSVAL(outbuf,10,0);  
  p = outbuf+12;
  strcpy(p,inbuf+12);  //right??
  p += NbtUtil::name_len(p);
  SSVAL(p,0,0x20);
  SSVAL(p,2,0x1);
  SIVAL(p,4,names_[n].ttl);
  SSVAL(p,8,6);
  SCVAL(p,10,nb_flags); //why??
  SCVAL(p,11,0);
  p += 12;

  memcpy(p,&myip_,4);      /* IP address of the name's owner (that's us) */
  p += 4;

  if (ip_equal(&ip,&bcastIp_))
  {
      Debug(0,"Not replying to broadcast address\n");
      return;
  }

  NbtUtil::show_nmb(outbuf);
  NbtUtil::send_packet(outbuf,NbtUtil::nmb_len(outbuf),&ip,myNetBiosPort,SOCK_DGRAM);

  return;
}

//
void NameService::replyNameQuery(char *inbuf, char *outbuf, const struct in_addr lastip)
{
  int name_trn_id = SVAL(inbuf,0);
  char qname[100]="";
  char *p = inbuf;
  unsigned char nb_flags = 0;
  struct in_addr tmpip;

  NbtUtil::name_extract(inbuf,12,qname);

  /* if it's not my name then return */
  if (!NbtUtil::name_equal(qname,myname_.c_str()))
    return;

  Debug(2,"(%s) querying my name (%s), sending positive reply\n",
	inet_ntoa(lastip), qname);
 
  /* Send a POSITIVE NAME QUERY RESPONSE to ack our presence */
  SSVAL(outbuf,0,name_trn_id);
  SCVAL(outbuf,2,((1<<7) | 0x5));
  SCVAL(outbuf,3,0);
  SSVAL(outbuf,4,0);
  SSVAL(outbuf,6,1);
  SSVAL(outbuf,8,0);
  SSVAL(outbuf,10,0);  
  p = outbuf+12;
  strcpy(p,inbuf+12);
  p += NbtUtil::name_len(p);
  SSVAL(p,0,0x20);
  SSVAL(p,2,0x1);
  SIVAL(p,4,myTTL_);
  SSVAL(p,8,6);
  SCVAL(p,10,nb_flags);
  SCVAL(p,11, 0);

  p += 12;
  memcpy(p,&myip_,4);
  p += 4;

  NbtUtil::show_nmb(outbuf);

  tmpip = lastip;
  NbtUtil::send_packet(outbuf,NbtUtil::nmb_len(outbuf),&tmpip,myNetBiosPort,SOCK_DGRAM);

  return;

}

//
void NameService::reply(char *inbuf, char *outbuf, const struct in_addr lastip)
{
    if (lastip.s_addr == myip_.s_addr)
        return;
  int opcode = CVAL(inbuf,2) >> 3;
  int nm_flags = ((CVAL(inbuf,2) & 0x7) << 4) + (CVAL(inbuf,3)>>4);
  int rcode = CVAL(inbuf,3) & 0xF;
  
  if (opcode == 0x5 && (nm_flags & ~1) == 0x10 && rcode == 0)
    replyRegRequest(inbuf,outbuf, lastip);
  if (opcode == 0x5 && (nm_flags & ~1) == 0x00 && rcode == 0)
	replyRegDemand(inbuf, outbuf, lastip);
  if (opcode == 0 && (nm_flags&~1) == 0x10 && rcode == 0)
    replyNameQuery(inbuf,outbuf, lastip);
     return;
}

void NameService::run()
{
  //注册本机名称
  bool res = registerName(myname_, &myip_); 
  if (!res)
  {
      Debug(0, "Failed to register my own name");
  }

  struct timeval prevTv, currTv;
  gettimeofday(&prevTv, 0);
  while (true)
  {
      struct in_addr lastip;
	  
	  
      if (!NbtUtil::receive_nmb(sockfd_, inBuffer_, 
        sizeof(inBuffer_), &lastip, 0 ))
          return;
      NbtUtil::show_nmb(inBuffer_);
      reply(inBuffer_, outBuffer_, lastip);
	  gettimeofday(&currTv, 0);
	  if (currTv.tv_sec - prevTv.tv_sec >= 30)
	  {
		  for (size_t k = 0; k < registerednames_.size(); ++k)
			  registerName(registerednames_[k], &myip_);
          gettimeofday(&prevTv, 0);
	  }

      Debug(2, "NameTables:\n");
      for (size_t k = 0; k < names_.size(); ++k)
      {
          if (names_[k].valid)
              Debug(2, "%s:%s\n", names_[k].name.c_str(), 
                      inet_ntoa(names_[k].ip));
      }

      ++nameTransId_;
  }
}

//find the namestr and return the index
int NameService::findName(string namestr)
{
    for (size_t i = 0; i < names_.size(); ++i)
    {
        if (names_[i].valid && NbtUtil::name_equal(
            names_[i].name.c_str(), namestr.c_str()))
            return i;
    }
    return -1;
}

int NameService::addName(string namestr, struct in_addr ip)
{
    name_t nameItem;
    nameItem.name = namestr;
    nameItem.valid = true;
    nameItem.ip.s_addr = ip.s_addr;
    names_.push_back(nameItem);
    return 0;
}
